import re
import matplotlib
import matplotlib.pyplot as plt
import os
datasink = False
root_dir = '/root/chenyp/ice_season/ice-predict/gc-ice-pred_obs/all_summary_lead6/2024_03_20_23_26_19_512_16/'
log_file_path = os.path.join(root_dir, 'train_2024_03_20_23_26_19_512_16.log')

loss_all, val_all = [], []
start_step = 0
fix_factor = 15.629847 / 274.69999



with open(log_file_path, 'r') as f:
    log_data = f.readlines()

for x in log_data:
    if x.startswith('epoch:'):
        loss_all.append(x)
    elif 'z500' in x or 'Z500' in x:
        val_all.append(x)
print(f"length of loss_all: {len(loss_all)}")
loss_list = []
print(loss_all[0])
for log in loss_all:
    x = re.split(":|,| ", log)
    loss_list.append(float(x[-1][:-2]))
# print(f"length of loss_list: {len(loss_list)}")
# print(loss_list[0])
val_list = []
for val in val_all:
    x = re.split(":|,| ", val)
    val_list.append([float(x[13]), float(x[17]), float(x[-5]), float(x[-1][:-2])])
print(f"length of val_list: {len(val_list)}")

def plt_log(data, title_name, x_label, y_label, loss=False, steps_per_epoch=113):
    # print(data)
    if loss:
        plt.plot([(i+1)/steps_per_epoch for i in range(len(data))], data, 'b', linewidth=2)
        plt.title(title_name, style='italic', fontsize=20)
    else:

        plt.plot([(i+1) for i in range(len(data))], data, 'b', linewidth=2)
        plt.title(f'{title_name}: the best epoch is {(data.index(min(data)) + 1 ) *5} with {round(min(data),2)}', style='italic', fontsize=12, loc='left')

    plt.ylabel(x_label, style='italic', fontsize=12)
    plt.xlabel(y_label, style='italic', fontsize=12)

fig = plt.figure(1, figsize=(12, 6))
if datasink:
    plt_log(loss_list, 'GraphCast Train Loss', 'Loss', 'epoch', loss=True)
else:
    # plt_log([data for data in loss_list if data < 0.001][::2000], 'GraphCast Train Loss', 'Loss', 'step', loss=True)
    plt_log(loss_list, 'GraphCast for ICE Pred Train Loss', 'Loss', 'epoch', loss=True)
plt.savefig(os.path.join(root_dir, "gc_ice_2_5e-4_step_113.png"))
# plt.savefig("./train_16-19_from_24pcs_150epochs/rank-10-device-2-23epochs.png")
# plt.close()

# fig2 = plt.figure(1, figsize=(18, 10))
# plt.subplot(4, 3, 1)
# plt_log([data[0] for data in val_list[::3]], 'Z500 6h', 'RMSE', '')
# plt.subplot(4, 3, 2)
# plt_log([data[0] for data in val_list[1::3]], 'Z500 72h', '', '')
# plt.subplot(4, 3, 3)
# plt_log([data[0] for data in val_list[2::3]], 'Z500 120h', '', '')

# plt.subplot(4, 3, 4)
# plt_log([data[1] for data in val_list[::3]], 'T850 6h', 'RMSE', '')
# plt.subplot(4, 3, 5)
# plt_log([data[1] for data in val_list[1::3]], 'T850 72h', '', '')
# plt.subplot(4, 3, 6)
# plt_log([data[1] for data in val_list[2::3]], 'T850 120h', '', '')

# plt.subplot(4, 3, 7)
# plt_log([data[2] for data in val_list[::3]], 'T2M 6h', 'RMSE', '')
# plt.subplot(4, 3, 8)
# plt_log([data[2] for data in val_list[1::3]], 'T2M 72h', '', '')
# plt.subplot(4, 3, 9)
# plt_log([data[2] for data in val_list[2::3]], 'T2M 120h', '', '')

# plt.subplot(4, 3, 10)
# plt_log([data[3] for data in val_list[::3]], 'U10 6h', 'RMSE', 'epoch')
# plt.subplot(4, 3, 11)
# plt_log([data[3] for data in val_list[1::3]], 'U10 72h', '', 'epoch')
# plt.subplot(4, 3, 12)
# plt_log([data[3] for data in val_list[2::3]], 'U10 120h', '', 'epoch')
# plt.subplots_adjust(wspace=0.25, hspace=0.6)
# plt.savefig("./rmses_tin_1_train_16-19.png")
# plt.show()
# plt.close()